# -*- encoding: utf-8 -*-
'''
@文件        :base_dao.py
@Description:       : 基础数据访问层
@Date     :2025/11/06 16:15:20
@Author      :lin
@version      :1.0
'''
from typing import Any, Dict, List, Optional, Type, Union, Tuple
from pydantic import BaseModel
from tortoise import Model, fields
from tortoise.exceptions import DoesNotExist, IntegrityError, OperationalError
from tortoise.transactions import in_transaction

# 自定义业务异常（可根据项目调整）
class DatabaseException(Exception):
    """数据库操作基础异常"""
    pass

class DataNotFoundError(DatabaseException):
    """数据不存在异常"""
    pass

class DataDuplicateError(DatabaseException):
    """数据重复异常（唯一索引冲突）"""
    pass

class TransactionError(DatabaseException):
    """事务执行异常"""
    pass

class BaseDAO:
    """
    Tortoise ORM 基础 DAO 类，封装通用 CRUD 操作
    所有实体 DAO 需继承此类，仅实现自身特殊查询逻辑
    """
    def __init__(self, model: Type[Model]):
        """
        初始化 DAO
        :param model: Tortoise ORM 模型类（如 User、Order）
        """
        self.model = model  # 绑定当前 DAO 操作的 ORM 模型

    # ------------------------------ 通用查询操作 ------------------------------
    async def get_by_id(self, pk: Any) -> Optional[Model]:
        """
        根据主键（ID）查询单条数据
        :param pk: 主键值（int/str 等，与模型主键类型一致）
        :return: 模型实例（不存在返回 None）
        """
        try:
            return await self.model.get(pk=pk)
        except DoesNotExist:
            return None
        except OperationalError as e:
            raise DatabaseException(f"查询数据失败：{str(e)}") from e

    async def get_one(
        self, 
        filters: Optional[Dict[str, Any]] = None,
        order_by: Optional[Union[str, List[str]]] = None
    ) -> Optional[Model]:
        """
        按条件查询单条数据（默认返回第一条）
        :param filters: 查询条件（如 {"username": "zhangsan", "status": 1}）
        :param order_by: 排序规则（如 "-create_time" 表示按创建时间倒序）
        :return: 模型实例（不存在返回 None）
        """
        query = self.model.filter(**(filters or {}))
        # 应用排序
        if order_by:
            query = query.order_by(*order_by) if isinstance(order_by, list) else query.order_by(order_by)
        try:
            return await query.first()
        except OperationalError as e:
            raise DatabaseException(f"查询数据失败：{str(e)}") from e

    async def get_list(
        self,
        filters: Optional[Dict[str, Any]] = None,
        order_by: Optional[Union[str, List[str]]] = None,
        limit: Optional[int] = None,
        offset: Optional[int] = None
    ) -> List[Model]:
        """
        按条件查询多条数据（支持排序、分页）
        :param filters: 查询条件
        :param order_by: 排序规则
        :param limit: 限制返回数量（默认不限制）
        :param offset: 偏移量（配合 limit 实现分页）
        :return: 模型实例列表
        """
        query = self.model.filter(**(filters or {}))
        # 排序
        if order_by:
            query = query.order_by(*order_by) if isinstance(order_by, list) else query.order_by(order_by)
        # 分页（offset + limit）
        if offset is not None:
            query = query.offset(offset)
        if limit is not None:
            query = query.limit(limit)
        try:
            return await query.all()
        except OperationalError as e:
            raise DatabaseException(f"查询数据列表失败：{str(e)}") from e

    async def get_paginated(
        self,
        page: int = 1,
        page_size: int = 10,
        filters: Optional[Dict[str, Any]] = None,
        order_by: Optional[Union[str, List[str]]] = None
    ) -> Dict[str, Any]:
        """
        分页查询（通用分页逻辑）
        :param page: 页码（默认第 1 页）
        :param page_size: 每页条数（默认 10 条）
        :param filters: 查询条件
        :param order_by: 排序规则
        :return: 分页结果（total: 总数, page: 当前页, page_size: 每页条数, total_pages: 总页数, items: 数据列表）
        """
        if page < 1:
            page = 1
        if page_size < 1:
            page_size = 10
        offset = (page - 1) * page_size

        # 1. 查询总数
        total = await self.model.filter(**(filters or {})).count()
        # 2. 查询当前页数据
        items = await self.get_list(
            filters=filters,
            order_by=order_by,
            limit=page_size,
            offset=offset
        )
        # 3. 计算总页数
        total_pages = (total + page_size - 1) // page_size

        return {
            "total": total,
            "page": page,
            "page_size": page_size,
            "total_pages": total_pages,
            "items": items
        }

    # ------------------------------ 通用写操作（新增/更新/删除） ------------------------------
    async def create(self, data: Union[Dict[str, Any], BaseModel]) -> Model:
        """
        新增数据（支持字典或 Pydantic 模型入参）
        :param data: 新增数据（字典或 Pydantic 模型实例）
        :return: 新增后的模型实例（包含数据库生成的主键、创建时间等字段）
        """
        # 处理 Pydantic 模型（转为字典）
        if isinstance(data, BaseModel):
            create_data = data.model_dump(exclude_unset=True)  # 只保留传入的字段
        else:
            create_data = data.copy()

        try:
            # Tortoise 模型 create 方法（自动提交）
            return await self.model.create(**create_data)
        except IntegrityError as e:
            # 捕获唯一索引冲突（如用户名重复）
            raise DataDuplicateError(f"数据已存在：{str(e)}") from e
        except OperationalError as e:
            raise DatabaseException(f"新增数据失败：{str(e)}") from e

    async def batch_create(self, data_list: List[Union[Dict[str, Any], BaseModel]]) -> List[Model]:
        """
        批量新增数据（高效批量插入，避免循环单条创建）
        :param data_list: 新增数据列表（元素为字典或 Pydantic 模型）
        :return: 新增后的模型实例列表
        """
        if not data_list:
            return []

        # 统一转为字典列表
        create_list = []
        for item in data_list:
            if isinstance(item, BaseModel):
                create_list.append(item.model_dump(exclude_unset=True))
            else:
                create_list.append(item.copy())

        try:
            # Tortoise 批量创建（内部优化为批量 SQL）
            return await self.model.bulk_create([self.model(**item) for item in create_list])
        except IntegrityError as e:
            raise DataDuplicateError(f"批量新增数据重复：{str(e)}") from e
        except OperationalError as e:
            raise DatabaseException(f"批量新增数据失败：{str(e)}") from e

    async def update(
        self,
        pk: Any,
        data: Union[Dict[str, Any], BaseModel],
        exclude_unset: bool = True
    ) -> Optional[Model]:
        """
        根据主键更新数据（支持部分更新）
        :param pk: 主键值
        :param data: 更新数据（字典或 Pydantic 模型）
        :param exclude_unset: 是否只更新传入的字段（True：部分更新，False：全量覆盖）
        :return: 更新后的模型实例（不存在返回 None）
        """
        # 1. 先查询数据是否存在
        db_obj = await self.get_by_id(pk)
        if not db_obj:
            return None

        # 2. 处理更新数据
        if isinstance(data, BaseModel):
            update_data = data.model_dump(exclude_unset=exclude_unset)
        else:
            update_data = data.copy()

        if not update_data:
            return db_obj  # 无更新字段，直接返回原数据

        try:
            # 3. 执行更新（Tortoise 模型的 update_from_dict 方法）
            await db_obj.update_from_dict(update_data)
            # 4. 提交更新（Tortoise 需要手动 save）
            await db_obj.save(update_fields=update_data.keys())
            return db_obj
        except IntegrityError as e:
            raise DataDuplicateError(f"更新后数据重复：{str(e)}") from e
        except OperationalError as e:
            raise DatabaseException(f"更新数据失败：{str(e)}") from e

    async def batch_update(
        self,
        filters: Dict[str, Any],
        update_data: Union[Dict[str, Any], BaseModel]
    ) -> int:
        """
        按条件批量更新数据
        :param filters: 更新条件（如 {"status": 0}）
        :param update_data: 更新内容（字典或 Pydantic 模型）
        :return: 受影响的行数
        """
        if not filters:
            raise DatabaseException("批量更新必须指定过滤条件（避免全表更新）")

        # 处理更新数据
        if isinstance(update_data, BaseModel):
            update_data = update_data.model_dump(exclude_unset=True)
        else:
            update_data = update_data.copy()

        if not update_data:
            return 0  # 无更新字段

        try:
            # Tortoise 批量更新（返回受影响行数）
            result = await self.model.filter(**filters).update(**update_data)
            return result
        except OperationalError as e:
            raise DatabaseException(f"批量更新数据失败：{str(e)}") from e

    async def delete(self, pk: Any) -> bool:
        """
        根据主键删除数据
        :param pk: 主键值
        :return: 删除成功返回 True，数据不存在返回 False
        """
        try:
            # 先查询数据是否存在
            db_obj = await self.get_by_id(pk)
            if not db_obj:
                return False
            # 执行删除
            await db_obj.delete()
            return True
        except OperationalError as e:
            raise DatabaseException(f"删除数据失败：{str(e)}") from e

    async def batch_delete(self, filters: Dict[str, Any]) -> int:
        """
        按条件批量删除数据
        :param filters: 删除条件（如 {"is_deleted": 1, "create_time__lt": "2024-01-01"}）
        :return: 被删除的行数
        """
        if not filters:
            raise DatabaseException("批量删除必须指定过滤条件（避免全表删除）")

        try:
            # Tortoise 批量删除（返回删除行数）
            deleted_count = await self.model.filter(**filters).delete()
            return deleted_count
        except OperationalError as e:
            raise DatabaseException(f"批量删除数据失败：{str(e)}") from e

    # ------------------------------ 事务操作 ------------------------------
    @staticmethod
    async def transaction(coroutines: List[Any]) -> List[Any]:
        """
        执行事务（多个数据库操作原子化，要么全成功，要么全回滚）
        :param coroutines: 事务中要执行的协程列表（如 [dao1.create(data1), dao2.update(pk2, data2)]）
        :return: 每个协程的执行结果列表
        """
        try:
            # 使用 Tortoise 内置事务上下文管理器
            async with in_transaction():
                results = []
                for coro in coroutines:
                    result = await coro
                    results.append(result)
                return results
        except Exception as e:
            # 事务执行失败，自动回滚，抛出统一异常
            raise TransactionError(f"事务执行失败：{str(e)}") from e